import jsonpickle
import openai
import os
import time

openai.api_key = os.getenv("OPENAI_API_KEY")

model_name = 'gpt-3.5-turbo-0613'

with open('./translator.prompt', 'r', encoding='utf-8') as f:
    prompt = f.read()


class ChatGpt(object):

    __init_interval = 2
    __max_internal_exponent = 5

    def __init__(self, model: str, messages: list, timeout: int = 30, max_retries: int = 3):
        self.model = model
        self.messages = messages
        self.timeout = timeout
        self.max_retries = max_retries

    def response(self):
        retry_time = 0
        # 进行指数退避重试
        while True:
            try:
                return openai.ChatCompletion.create(
                    model=self.model,
                    messages=self.messages,
                    timeout=self.timeout
                )
            except openai.error.RateLimitError as e:
                interval = self.__init_interval * \
                    (1 << min(self.__max_internal_exponent, retry_time))
                print(
                    f"[warning]: Slepp {interval}s and retry due to rate limitation")
                time.sleep(interval)
                retry_time += 1


def translate(content: str):
    messages = [{'role': 'system', 'content': prompt}]
    messages.append({'role': 'assistant', "content": 'yes'})
    messages.append({'role': 'user', 'content': jsonpickle.encode({
        "source": content
    }, indent=True)})
    gpt = ChatGpt(model=model_name, messages=messages)
    gpt_resp = gpt.response()
    if gpt_resp.choices and gpt_resp.choices[0]:
        resp = gpt_resp.choices[0].message.content
        resp_obj = jsonpickle.decode(resp)
        if True is resp_obj['success']:
            return resp_obj['translation']
        else:
            if False is resp_obj['success'] and resp_obj['reason']:
                raise Exception(resp_obj['reason'])
    return None


if __name__ == '__main__':
    while True:
        source = input("Please input, enter Q to quit: ")
        if source == 'Q':
            print('Bye!')
            break
        print(f">>> {translate(source)}")
